import torch
from sklearn.cluster import KMeans
from sklearn.cluster import DBSCAN
from sklearn.mixture import GaussianMixture
import numpy as np
import math

def generate_y_threshold(graph_data, threshold=0.9):
    """
    方法 1：阈值法 - 计算所有节点的整体风险得分，并统一赋值 y 标签。
    数量少的类标记为 1（脆弱节点），数量多的类标记为 0（非脆弱节点）。
    """
    risk_scores = graph_data.x[:, :65].sum(dim=1).float()  # 计算所有节点的风险得分
    y_labels = (risk_scores > threshold).long()  # 根据阈值划分 0/1

    # 统计每个标签数量
    # counts = torch.bincount(y_labels)

    # 找到数量少的那一类
    # if counts[0] <= counts[1]:
    #     # 0 类是脆弱节点，需要翻转标签
    #     y_labels = 1 - y_labels
    # 如果 1 类本来就少，就不用动

    graph_data.y = y_labels  # 统一赋值 y
    return graph_data

def generate_y_kmeans(graph_data, n_clusters=2):
    """
    方法 2：K-Means 聚类法 - 先考虑所有节点特征，再统一赋值 y 标签。
    """
    X_numpy = graph_data.x[:, :65].cpu().numpy()  # 取前 65 维特征
    kmeans = KMeans(n_clusters=n_clusters, random_state=42)
    kmeans.fit(X_numpy)

    y_labels = torch.tensor(kmeans.labels_, dtype=torch.long)  # 统一计算所有 y

    # 统计每个类别数量
    counts = torch.bincount(y_labels)

    # 找到数量少的那个类别
    if counts[0] <= counts[1]:
        # 0 类是脆弱节点，需要把 0 -> 1, 1 -> 0
        y_labels = 1 - y_labels
    # 如果 1 类本身就少，就不用动

    graph_data.y = y_labels
    return graph_data


def generate_y_dbscan(graph_data, eps=0.5, min_samples=5):
    """
    方法 3：DBSCAN 聚类法 - 基于密度，多提取中间11维图结构，再统一赋值 y 标签。
    """
    X_numpy = graph_data.x[:, :76].cpu().numpy() # 取前 76 维特征
    dbscan = DBSCAN(eps=eps, min_samples=min_samples)
    labels = dbscan.fit_predict(X_numpy)
    graph_data.y = torch.tensor(labels, dtype=torch.long)
    return graph_data

def generate_y_gmm(graph_data, n_components=2):
    """
    方法 4：GMM - 高斯分布概率模型，多加入最后一维，再统一赋值 y 标签。
    """
    #X_numpy = graph_data.x[:, :65].cpu().numpy() # 取前 65 维特征
    X_numpy = np.concatenate([
        graph_data.x[:, :65].cpu().numpy(),  # 高危端口和服务
        graph_data.x[:, [-1]].cpu().numpy()  # OS指数
    ], axis=1)

    gmm = GaussianMixture(n_components=n_components, random_state=42)
    gmm.fit(X_numpy)  # 特征拟合GMM模型
    labels = gmm.predict(X_numpy)  # 预测样本标签

    labels_tensor = torch.tensor(labels, dtype=torch.long)

    # 统计每个类别数量
    counts = torch.bincount(labels_tensor)

    # 找到数量少的那个类别
    if counts[0] <= counts[1]:
        # 0 类是脆弱节点，需要翻转标签
        labels_tensor = 1 - labels_tensor
    # 如果 1 类本来就少，就不用动

    graph_data.y = labels_tensor
    return graph_data

def generate_y_combined(graph_data, k=1.0, hub_dom_threshold=0.1, bridge_risk_threshold=0.2, spread_risk_threshold=0.1,
                        transmission_efficiency_threshold=3.0, zoom_levels=None):
    """
    自动化标签生成函数 - 结合第一层和第二层维度进行组合脆弱性评估，并赋值 y 标签。

    参数:
    - graph_data: PyTorch Geometric Data 对象
    - k: 常数系数，用于 Risk_Portfolio 计算
    - hub_dom_threshold: Hub 节点支配性的阈值
    - bridge_risk_threshold: 桥边关键性的阈值
    - spread_risk_threshold: 攻击扩散关键性的阈值
    - transmission_efficiency_threshold: 脆弱性传递效率的阈值

    返回:
    - 更新后的 graph_data，包含 y 标签

    TODO.Katz统计为0需要调整
    """
    # 提取节点特征
    x = graph_data.x
    risk_port_count = x[:, -3]  # 风险端口数
    risk_service_count = x[:, -2]  # 风险服务数
    os_security_index = x[:, -1]  # OS安全指数
    degree = x[:, 65]  # 度数
    eigenvector_centrality = x[:, 72]  # 特征向量中心性
    betweenness_centrality = x[:, 69]  # 介数中心性
    katz_centrality = x[:, 75]  # Katz中心性

    # 计算第二层节点维度
    hub_dom = torch.log(degree + 1) * eigenvector_centrality
    spread_risk = betweenness_centrality * (katz_centrality ** 2)
    # 变焦修正 risk_portfolio
    if zoom_levels is not None:
        rp_correction = 1 - 0.5 / (1 + (zoom_levels/2.5)**4)
    else:
        rp_correction = 0
    risk_portfolio = torch.tanh(risk_port_count * 2 / 7) * risk_service_count - k * (3 + rp_correction - os_security_index)

    # 提取边特征
    edge_attr = graph_data.edge_attr
    is_bridge = edge_attr[:, 5]  # 是否为桥边
    centrality_diff = edge_attr[:, 2]  # 中心性差
    weight = edge_attr[:, 0]  # 权重
    vulnerability_transmission = edge_attr[:, 4]  # 脆弱性传递概率

    # 计算第二层边维度
    bridge_risk = is_bridge * centrality_diff
    transmission_efficiency = weight * vulnerability_transmission

    # 变焦自适应阈值 hub_dom_threshold=ln((e^0.1-1)) * 2^x + 1)
    if zoom_levels is not None:
        hub_dom_threshold *= 10 * math.log((math.exp(0.1) - 1) * (2 ** zoom_levels) + 1)

    # 组合脆弱性评估
    node_vulnerability_scores = (
            (hub_dom > hub_dom_threshold).float() +
            (spread_risk > spread_risk_threshold).float() +
            (risk_portfolio > 0).float()  # Risk_Portfolio < 0 表示高风险
    )

    # 边的脆弱性评估（这里仅作为参考，不直接影响节点标签）
    edge_vulnerability_scores = (
            (bridge_risk > bridge_risk_threshold).float() +
            (transmission_efficiency > transmission_efficiency_threshold).float()
    )

    # 赋予 y 标签
    y_labels = (node_vulnerability_scores >= 1).long()  # 至少满足一个条件则标记为脆弱节点
    graph_data.y = y_labels

    return graph_data


